"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$HeadURL$

Created on:           Apr 21, 2010
Last Revised on:      $Date$
Version:              $Revision$
ID:                   $Id$
Author:               $Author$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================

"""

#===============================================================================
# Imports
#===============================================================================
import JT
import vtk
import numpy as np

class Transform(vtk.vtkTransform):
    '''
    classdocs
    '''


    def __init__(self):
        '''
        Constructor
        '''
        # in the real world scenario, this pose would not be known, and would 
        # Therefore be set to the identity
        # But for Gold Standard data, it would be given
        self.volume_transform = vtk.vtkTransform()
        self.volume_transform.PostMultiply()
        
        # The body center is the offset of the center of rotation from the 
        # the origin of the volume
        self.body_center = vtk.vtkTransform()
        self.body_center.PostMultiply()
        self.body_center_inv = vtk.vtkTransform()
        self.body_center_inv.PostMultiply()
        
        # in  the real world scenario, this pose is not used, and would be 
        # set to the identity
        # For Gold Standard data, this is the standardized starting position.
        self.start_error = vtk.vtkTransform()
        self.start_error.PostMultiply()
        
        # Set the behavior to post-multiply
        self.PostMultiply()
        
        self.body_center_inv.SetInput(self.volume_transform)
        self.start_error.SetInput(self.body_center_inv)
        self.body_center.SetInput(self.start_error)
        self.SetInput(self.body_center)
        
    def set_volume_transform(self,transform):
        """
        Transform is a 4x4 numpy array.
        Set the known Gold Standard transform of the volume.  This will be 
        the identity in real world scenarios.
        """
        mat = vtk.vtkMatrix4x4()
        for row in range(4):
            for col in range(4):
                mat.SetElement(row,col,transform[row][col])
        
        self.volume_transform.Identity()
        self.volume_transform.Concatenate(mat)
        
        
    def set_body_center(self,center):
        """
        center is a 3 element array representing the offset of the body 
        center of roatation relative to the volume origin.  
        """
        mat1 = vtk.vtkMatrix4x4()
        
        mat1.SetElement(0 , 3, center[0])
        mat1.SetElement(1 , 3, center[1])
        mat1.SetElement(2 , 3, center[2])
        
        self.body_center.Identity()
        self.body_center.Concatenate(mat1)
        
        mat2 = vtk.vtkMatrix4x4()
        
        mat2.SetElement(0 , 3, -center[0])
        mat2.SetElement(1 , 3, -center[1])
        mat2.SetElement(2 , 3, -center[2])
        
        self.body_center_inv.Identity()
        self.body_center_inv.Concatenate(mat2)

        
    def set_start_error(self,pose):
        """
        pose is a 6 element vector (Tx,Ty,Tz,Rx,Ry,Rz) 
        representing the known offset of the volume from the true position
        Taken from start positions for Gold Standard data
        """
        self.set_pose(pose)
    
    def set_pose(self,pose):
        """
        pose is a 6 element vector (Tx,Ty,Tz,Rx,Ry,Rz) 
        That describes the position and orientation of the volume.
        """
        transform = self.get_transform_from_pose(pose)
        # if the elements are set directly, they will get overwritten when the
        # pipeline updates, concatenate a new matrix onto the pipeline.
        mat = vtk.vtkMatrix4x4()
        for row in range(4):
            for col in range(4):
                mat.SetElement(row,col,transform[row][col])
        
        self.start_error.Identity()
        self.start_error.Concatenate(mat)
    
    def get_pose(self):
        """
        Return the current 6 element vector (Tx,Ty,Tz,Rx,Ry,Rz)
        That describes the position and orientation of the volume.
        """
        vtkMatrix = self.GetMatrix()
        transform = JT.utils.GetNumpyArrayFromVTKMatrix(vtkMatrix)
        pose = self.get_pose_from_transform(transform)
        return pose
        
    def get_transform_from_pose(self,pose):
        """
        Given an arbitrary pose (6 element vector), return a transform (4x4 numpy array)
        derived from T=Rz*Ry*Rx.
        """
        Tx,Ty,Tz,Rx,Ry,Rz = pose
        
        # transform is TStartTransformation (body center to start pose)
        transform = np.eye(4)
        transform[0,0] = np.cos(Ry)*np.cos(Rz)
        transform[0,1] = -np.cos(Rx)*np.sin(Rz) + np.cos(Rz)*np.sin(Rx)*np.sin(Ry)
        transform[0,2] = np.sin(Rx)*np.sin(Rz) + np.cos(Rx)*np.cos(Rz)*np.sin(Ry)
        transform[0,3] = Tx
        transform[1,0] = np.cos(Ry)*np.sin(Rz)
        transform[1,1] = np.cos(Rx)*np.cos(Rz) + np.sin(Rx)*np.sin(Ry)*np.sin(Rz)
        transform[1,2] = -np.cos(Rz)*np.sin(Rx) + np.cos(Rx)*np.sin(Ry)*np.sin(Rz)
        transform[1,3] = Ty
        transform[2,0] = -np.sin(Ry)
        transform[2,1] = np.cos(Ry)*np.sin(Rx)
        transform[2,2] = np.cos(Rx)*np.cos(Ry)
        transform[2,3] = Tz
        
        return transform 
        
    def get_pose_from_transform(self,transform):
        """
        Given a 4x4 transform in numpy, return a 6 element vector 
        (Tx,Ty,Tz,Rx,Ry,Rz).  Rotation order: T=Rz*Ry*Rx [x->y->z]
        rotations are in radians.
        """
        Tx = transform[0][3]
        Ty = transform[1][3]
        Tz = transform[2][3]
        Ry = -np.arcsin(transform[2][0])
        Rx = np.arcsin(transform[2][1]*1/np.cos(Ry))
        Rz = np.arcsin(transform[1][0]*1/np.cos(Ry)) 
        
        pose=np.array([Tx,Ty,Tz,Rx,Ry,Rz])
        
        return pose
        